﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.ComponentModel;
using System.Configuration.Install;
using System.Data;
using System.Diagnostics;
using System.Drawing;
using System.Linq;
using System.Runtime.InteropServices;
using System.Security;
using System.ServiceProcess;
using System.Text;
using System.Windows.Forms;
using ServiceUtilities;

namespace ServiceUtilities
{
    public static class ServiceControllerExtensions
    {
        public static void StartOnInstall(this ServiceInstaller installer)
        {
            installer.AfterInstall += (o, e) => StartService(installer.ServiceName, new TimeSpan(0, 0, 30));
        }

        public static void StartService(string serviceName, TimeSpan waitTo)
        {
            // Get a handle to the service
            using (var sc = new ServiceController(serviceName))
            {
                sc.Start();
                sc.WaitForStatus(ServiceControllerStatus.Running, waitTo);
                sc.Close();
            }
        }

        public static void SetRecoveryOptions(
            this ServiceInstaller installer,
            TimeSpan resetAfter,
            string commandToExecuteOnFail,
            string rebootMessage,
            params FailureAction[] failureActions)
        {
            installer.AfterInstall +=
                (o, e) => WithServiceHandle(installer.ServiceName,
                                            handle => SetRecovery(handle, resetAfter, commandToExecuteOnFail, rebootMessage, failureActions));
        }

        private static void WithServiceHandle(string serviceName, Action<IntPtr> useHandle)
        {
            var scmHndl = Win32.OpenSCManager(null, null, Win32.SC_MANAGER_ALL_ACCESS);
            if (scmHndl.ToInt32() <= 0)
                throw new Exception("Failed to Open Service Control Manager");
            try
            {
                var svcLock = Win32.LockServiceDatabase(scmHndl);

                if (svcLock.ToInt32() <= 0)
                    throw new Exception("Failed to Lock Service Database for Write");

                try
                {
                    var svcHndl = Win32.OpenService(scmHndl, serviceName, Win32.SERVICE_ALL_ACCESS);
                    if (svcHndl.ToInt32() <= 0)
                        throw new Exception("Failed to Open Service ");
                    try
                    {
                        useHandle(svcHndl);
                    }
                    finally
                    {
                        Win32.CloseServiceHandle(scmHndl);
                    }
                }
                finally
                {
                    Win32.UnlockServiceDatabase(svcLock);
                }
            }
            finally
            {
                Win32.UnlockServiceDatabase(scmHndl);
            }
        }


        private static void SetRecovery(
            IntPtr svcHndl,
            TimeSpan resetAfter,
            string commandToExecuteOnFail,
            string rebootMessage,
            params FailureAction[] failureActions
            )
        {
            // Circumvent windows logic of making actions 2 & 3 mimmick the first if not specified.
            if (failureActions == null || failureActions.Length == 0)
                failureActions = new[] { FailureAction.None };
            else
                if (failureActions.Length < 3)
                    failureActions = failureActions.Concat(new[] {FailureAction.None}).ToArray();

            // If we need shutdown privilege, then grant it to this process
            if (failureActions.Any(fa => fa.Type == RecoverAction.Reboot))
                GrantShutdownPrivilege();

            var actions = failureActions.SelectMany(fa => new[] { (int)fa.Type, fa.Delay }).ToArray();
            var actionsIsInts = Marshal.AllocHGlobal(actions.Length * sizeof(int));
            try
            {
                Marshal.Copy(actions, 0, actionsIsInts, actions.Length);

                // Set the SERVICE_FAILURE_ACTIONS struct
                var sfa = new ServiceFailureActions
                              {
                                  dwResetPeriod =
                                      resetAfter == TimeSpan.Zero
                                          ? Win32.SERVICE_NO_CHANGE
                                          : (int)resetAfter.TotalSeconds,
                                  cActions = failureActions.Length,
                                  lpCommand = commandToExecuteOnFail,
                                  lpRebootMsg = rebootMessage,
                                  lpsaActions = actionsIsInts.ToInt32()
                              };

                if (Win32.ChangeServiceFailureActions(svcHndl, Win32.SERVICE_CONFIG_FAILURE_ACTIONS, ref sfa))
                {
                    var err = Win32.GetLastError();
                    if (err == Win32.ERROR_ACCESS_DENIED)
                        throw new SecurityException("Access Denied while setting Failure Actions");
                }
            }
            finally
            {
                Marshal.FreeHGlobal(actionsIsInts);
            }
        }

        private static void GrantShutdownPrivilege()
        {
            IntPtr hToken = IntPtr.Zero;

            var myProc = Win32.GetCurrentProcess();

            if (!Win32.OpenProcessToken(myProc, Win32.TOKEN_ADJUST_PRIVILEGES | Win32.TOKEN_QUERY, ref hToken))
                throw new SecurityException("Failed to grant shutdown privilege");

            try
            {
                long luid = 0;
                Win32.LookupPrivilegeValue(null, Win32.SE_SHUTDOWN_NAME, ref luid);

                var tkp = new TokenPrivileges
                              {
                                  PrivilegeCount = 1,
                                  Privileges = { Luid = luid, Attributes = Win32.SE_PRIVILEGE_ENABLED }
                              };

                int retLen = 0;
                if (!Win32.AdjustTokenPrivileges(hToken, false, ref tkp, 0, IntPtr.Zero, ref retLen)
                    && Win32.GetLastError() != 0)
                    throw new SecurityException("Failed to grant shutdown privilege");
            }
            finally
            {
                Win32.CloseHandle(hToken);
            }
        }
    }
}